from wwengine.datasets import build_dataset
from wwengine.config import Config
from wwengine.utils import get_draw_box, get_arrows
import open3d as o3d

if __name__ == '__main__':
    cfg = Config.fromfile('files/default.json')
    tasks = cfg['TASKS']
    hw_data_train = build_dataset(cfg['DATASETS'], tasks, split='train')
    hw_data_val = build_dataset(cfg['DATASETS'], tasks, split='val')
    hw_data_test = build_dataset(cfg['DATASETS'], tasks, split='test')
    data_dict = hw_data_val.__getitem__(3)

    # 可视化
    vis = o3d.visualization.Visualizer()
    vis.create_window(window_name=data_dict['dataset_name'] + data_dict['data_name'])
    render_option = vis.get_render_option()
    render_option.point_size = 2
    coordinate_frame = o3d.geometry.TriangleMesh.create_coordinate_frame(size=5.0, origin=[0, 0, 0])
    vis.add_geometry(coordinate_frame)

    if data_dict['lidar_pts'] is not None:
        data_lidar = o3d.geometry.PointCloud()
        data_lidar.points = o3d.utility.Vector3dVector(data_dict['lidar_pts'])
        # 可视化激光点云
        data_lidar.paint_uniform_color([0.3, 0.5, 0])
        if 'SEG' in tasks:  # 分割任务才显示点云颜色
            if data_dict['gt_seg'] is not None:
                gt_seg = data_dict['gt_seg']
                for pt_id in range(gt_seg.shape[0]):
                    sem = gt_seg[pt_id][0]
                    if sem < 20:
                        data_lidar.colors[pt_id] = [0.05 * sem, 0, 0]
                    elif sem >= 20 and sem < 40:
                        data_lidar.colors[pt_id] = [0, 0.05 * sem, 0]
                    else:
                        data_lidar.colors[pt_id] = [0, 0, 0.05 * sem]
        vis.add_geometry(data_lidar)
    # # 可视化毫米波点云 注意是否数据增强
    if data_dict['radar_pts'] is not None:
        data_radar = o3d.geometry.PointCloud()
        data_radar.points = o3d.utility.Vector3dVector(data_dict['radar_pts'][:, :3])
        data_radar.paint_uniform_color([1, 0.706, 0.1])
        render_option.point_size = 5.0  # 设置渲染点的大小
        vis.add_geometry(data_radar)
    # 可视化boxes
    if data_dict['gt_bboxes_3d'] is not None:
        boxes = data_dict['gt_bboxes_3d']
        exp_draw_boxes = get_draw_box(boxes)
        arrows = get_arrows(boxes)
        box_id = -1
        for box in exp_draw_boxes:
            box_id = box_id + 1
            if 'old_bboxes_size' in data_dict.keys():
                if box_id >= data_dict['old_bboxes_size']:
                    box.paint_uniform_color([0.4, 0.2, 0.9])  # 数据增强的box换色
            vis.add_geometry(box)
        for arrow in arrows:
            vis.add_geometry(arrow)

    vis.run()
    vis.destroy_window()
